from typing import Dict
from pathlib import Path
import logging
import torch
from transformers import SeamlessM4TProcessor, SeamlessM4TModel
from core.translation_interface import TranslationEngine

logger = logging.getLogger(__name__)

class SeamlessM4TAdapter(TranslationEngine):
    """
    SeamlessM4T 端到端模型的适配器
    直接实现音频到目标语言的翻译
    """
    
    def __init__(self, model_name: str = "facebook/seamless-m4t-medium"):
        self.model_name = model_name
        self.processor = None
        self.model = None
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self._load_model()
        
    def _load_model(self):
        """加载SeamlessM4T模型"""
        try:
            logger.info(f"正在加载SeamlessM4T模型: {self.model_name}")
            self.processor = SeamlessM4TProcessor.from_pretrained(self.model_name)
            self.model = SeamlessM4TModel.from_pretrained(self.model_name)
            self.model.to(self.device)
            logger.info("SeamlessM4T模型加载完成")
        except Exception as e:
            logger.error(f"SeamlessM4T模型加载失败: {e}")
            raise
    
    def translate_audio(
        self, 
        audio_path: str, 
        source_lang: str, 
        target_lang: str,
        enable_translation: bool = True,
        return_segments: bool = False
    ) -> Dict:
        """端到端音频翻译"""
        
        if not self.is_ready():
            raise RuntimeError("SeamlessM4T模型未就绪")
            
        try:
            # 加载音频
            audio_input = self.processor(
                audio=audio_path, 
                return_tensors="pt"
            ).to(self.device)
            
            # 设置语言代码
            source_lang_code = self._get_lang_code(source_lang)
            target_lang_code = self._get_lang_code(target_lang)
            
            # 生成翻译
            if enable_translation:
                # 音频 -> 目标语言文本
                output_tokens = self.model.generate(
                    **audio_input,
                    tgt_lang=target_lang_code,
                    src_lang=source_lang_code,
                    max_new_tokens=256
                )
                translated_text = self.processor.decode(
                    output_tokens[0].tolist(), 
                    skip_special_tokens=True
                )
                
                # 获取源语言文本（用于显示）
                source_tokens = self.model.generate(
                    **audio_input,
                    tgt_lang=source_lang_code,
                    src_lang=source_lang_code,
                    max_new_tokens=256
                )
                source_text = self.processor.decode(
                    source_tokens[0].tolist(), 
                    skip_special_tokens=True
                )
            else:
                # 只做ASR
                source_tokens = self.model.generate(
                    **audio_input,
                    tgt_lang=source_lang_code,
                    src_lang=source_lang_code,
                    max_new_tokens=256
                )
                source_text = self.processor.decode(
                    source_tokens[0].tolist(), 
                    skip_special_tokens=True
                )
                translated_text = source_text
            
            return {
                "text": source_text,
                "translated_text": translated_text,
                "confidence": 0.95,  # SeamlessM4T通常有较高的置信度
                "duration": 0.0,  # 需要单独计算
                "segments": [] if not return_segments else self._extract_segments(audio_input)
            }
            
        except Exception as e:
            logger.error(f"SeamlessM4T翻译失败: {e}")
            raise
    
    def _get_lang_code(self, lang: str) -> str:
        """将语言代码转换为SeamlessM4T格式"""
        lang_mapping = {
            "zh": "cmn", "en": "eng", "ja": "jpn", "ko": "kor",
            "fr": "fra", "de": "deu", "es": "spa", "ru": "rus"
        }
        return lang_mapping.get(lang, "eng")
    
    def _extract_segments(self, audio_input) -> list:
        """提取音频分段信息（简化实现）"""
        # 这里可以实现更复杂的分段逻辑
        return []
    
    def is_ready(self) -> bool:
        """检查模型是否就绪"""
        return self.processor is not None and self.model is not None
    
    def get_supported_languages(self) -> Dict[str, list]:
        """获取支持的语言列表"""
        return {
            "asr_languages": ["zh", "en", "ja", "ko", "fr", "de", "es", "ru"],
            "translation_languages": ["zh", "en", "ja", "ko", "fr", "de", "es", "ru"]
        }




